from reid_libs.utils.testing_utils.img_evaluators import ImgEvaluator
from reid_libs.datasets_manager import get_dataset
from reid_libs.utils.data_utils.person_transform import personTestTransform
from reid_libs.utils.data_utils.person_dataloader import PersonImageDataset
from torch.utils.data import DataLoader
from reid_libs.models import create_model
from IPython import embed
import torch

# 测试阶段-选择设备
# device = torch.device("cuda")
device = torch.device("cpu")

# 测试阶段-超参数
BATCH_SIZE = 1
NUM_WORKERS = 0

# 测试的数据集
Test_Dataset = get_dataset('market1501', r'D:\reid\DATA4REID')

query_loader = DataLoader(
    PersonImageDataset(Test_Dataset.query, transform=personTestTransform),
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
    shuffle=False,
    pin_memory=True)

gallery_loader = DataLoader(
    PersonImageDataset(Test_Dataset.gallery, transform=personTestTransform),
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
    shuffle=False,
    pin_memory=True)

# 测试的模型
model = create_model(name='resnet50_rga',
                     num_classes=751,
                     height=256,
                     width=128,
                     pretrained=True,
                     num_feat=2048)


# # 并行化
# model = torch.nn.DataParallel(model).to(device)


# 载入模型权重
resume = r'good_ckpt\market1501_resnet50rga_b64n4f2048\checkpoint_200.pth.tar'
print("Loading checkpoint from '{}'".format(resume))
checkpoint = torch.load(resume, map_location='cpu')
model.load_state_dict(checkpoint['state_dict'])
test_epoch = checkpoint['epoch']
print('test epoch:', test_epoch)


# 并行化
model = torch.nn.DataParallel(model).to(device)


# test/evaluate the model
evaluator = ImgEvaluator(model, selected_device=device)
# feat_ : 未归一化 ; feat : 归一化后
# feats_list = ['feat_', 'feat']
feats_list = ['feat_', 'feat']


evaluator.eval_worerank(query_loader,
                        gallery_loader,
                        Test_Dataset.query,
                        Test_Dataset.gallery,
                        metric=['euclidean', 'cosine'],
                        types_list=feats_list)
